import functools
import re
from copy import deepcopy
from typing import Optional, Dict

from langchain_core.messages import AIMessage, ChatMessage, HumanMessage
from docx import Document

from lc_agents import section_divider_agent, check_agent_exec, writer, write_history_message
from utils.retry import retry
import typer
from typing_extensions import Annotated
import rich


def parse_section(ai_message: AIMessage) -> dict:
    """Parse the AI message."""
    content = ai_message.content
    res = re.findall("```json([\s\S]*?)```[\s]*", content)[0]
    return eval(res)


def parse_word_count(ai_message: dict) -> int:
    try:
        content = ai_message['output']
        res = re.findall("\d+", content)[0]
        return int(res)
    except Exception as e:
        return 0


@retry(tries=3)
def write_section(intput: str, content: str) -> dict:
    message = section_divider_agent.invoke({"input": intput, "content": content})
    sections = parse_section(message)
    return sections


@retry(tries=3)
def check_num(input) -> int:
    res = check_agent_exec.invoke({"input": input})
    return parse_word_count(res)


def calculate_word_for_section(section_keys: dict, total: int) -> int:
    cnt = 0
    for h1, v1 in section_keys.items():
        for h2, v2 in v1.items():
            if len(v2) == 0:
                cnt += 1
            else:
                cnt += len(v2)

    return (total // cnt) + 10


def main(title: Annotated[str, typer.Option()], content: Annotated[str, typer.Option()], total_words: int = 10000,
         human_section: Annotated[Optional[str], typer.Option(help="人为设定标题，dict格式")] = None,
         output_file: Annotated[str, typer.Option()] = 'output.docx', ):
    document = Document()
    try:
        if human_section is None:
            section_out = section_divider_agent.invoke({"input": title,
                                                        "content": content})
            section_keys = parse_section(section_out)
        else:
            section_keys = eval(human_section)
        min_section_word = calculate_word_for_section(section_keys, total_words)
        rich.print(f"[green]SYSTEM: 一共需要书写 {total_words}，每个section需要撰写{min_section_word}. [/green]")
        document.add_heading(title, 0)
    except Exception as e:
        rich.print(f"[red]标题解析错误: {e}")
        exit(1)
    try:
        for h1, v1 in section_keys.items():
            document.add_heading(h1, 1)
            for h2, v2 in v1.items():
                document.add_heading(h2, 2)
                if len(v2) > 0:
                    for h3 in v2:
                        document.add_heading(h3, 3)
                        memory = llm_write_section(content, h3, title, min_section_word)
                        document.add_paragraph('\n'.join(memory))
                else:
                    memory = llm_write_section(content, h2, title, min_section_word)
                    document.add_paragraph('\n'.join(memory))

        document.save(f'{output_file}.docx')
    except Exception as e:
        rich.print(f"[red] {e} [/red]",)
        document.save(f'{output_file}.docx')


def llm_write_section(content, head, title, word_count: int):
    memory = []
    history = deepcopy(write_history_message)
    history.append(
        HumanMessage(content=f"我想要写一篇{title}的文章，内容包括{content}，现在请你帮我撰写 {head}")
    )
    message = writer.invoke(history)
    history.append(message)
    output = message.content
    memory.append(output)
    cnt = 0
    check_out = check_agent_exec.invoke({"input": output})
    cnt += parse_word_count(check_out)
    while cnt < word_count:
        message = writer.invoke(
            [*history, HumanMessage(content=f"请你继续撰写")]
        )
        output = message.content
        memory.append(output)
        check_out = check_agent_exec.invoke({"input": output})
        cnt = cnt + parse_word_count(check_out)
    return memory


if __name__ == '__main__':
    # main(title="基于Vus.js的美食分享管理系统的开发与实践。",
    #      content="美食的收藏、分享、浏览评论功能、管理员：用户管理、美食管理、评论管理")
    # print(calculate_word_for_section({"h1":{"h2":["h3","h3"],"h2_2":["a",'b'],"h3":[]}},5000))
    typer.run(main)
